import numpy as np
import scipy.io.wavfile as wf
import sklearn.svm as svm
import sklearn.metrics as sm
import sklearn.preprocessing as sp
import os
import pickle
from scipy.fftpack import fft

# 分析音频文件中的主要频率
def analyze_frequency(wav_filename):
    sample_rate, data = wf.read(wav_filename)
    if data.ndim > 1:
        data = data[:, 0]
    fft_out = fft(data)
    freqs = np.fft.fftfreq(len(fft_out))
    idx = np.argmax(np.abs(fft_out))
    freq = freqs[idx]
    freq_in_hertz = abs(freq * sample_rate)
    return freq_in_hertz

# 检索目录下的所有wav文件
def search_files(directory):
    files_dict = {}
    for cur_dir, sub_dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".wav"):
                label = cur_dir.split(os.path.sep)[-1]
                if label not in files_dict:
                    files_dict[label] = []
                url = os.path.join(cur_dir, file)
                files_dict[label].append(url)
    return files_dict

# 提取音频文件的最大幅度频率
def files_frequency(file_urls):
    x_data, y_data = [], []
    for label, urls in file_urls.items():
        for file in urls:
            frequency = analyze_frequency(file)
            x_data.append([frequency])
            y_data.append(label)
    x_data = np.array(x_data)
    return x_data, y_data

# 读取训练集
train_urls = search_files("D:/Pycharm/trainvoice_deployment/train")
train_x, train_y = files_frequency(train_urls)

# 确保train_x不为空
if train_x.size == 0:
    print("No data found. Please check the file paths and data extraction process.")
else:
    # 如果数据只有一个特征，重塑train_x为二维数组
    train_x = train_x.reshape(-1, 1)

    # 标签编码
    encoder = sp.LabelEncoder()
    train_y_label = encoder.fit_transform(train_y)

    # 创建SVM模型并训练
    model = svm.SVC(kernel="poly", degree=2, gamma="auto", probability=True)
    model.fit(train_x, train_y_label)


    # # # 保存模型：已经训练好了，别再保存了
    # # with open(r"D:\graduation design\model_deployment\voc_svm\voc_svm\model_mfc.pickle", "wb") as file:
    #     pickle.dump(model, file)
